import numpy as np
from abc import ABC, abstractmethod
from collections import defaultdict


class BaseRewardFunction(ABC):
    """
    基类奖励函数
    Reward-specific reset and get_reward methods在子类中实现
    """
    def __init__(self, config):
        self.config = config
        # 内部变量
        self.reward_scale = getattr(self.config, f'{self.__class__.__name__}_scale', 1.0) # 奖励缩放因子
        self.is_potential = getattr(self.config, f'{self.__class__.__name__}_potential', False) # 是否为潜在奖励函数
        self.pre_rewards = defaultdict(float) # 存储之前的奖励
        self.reward_trajectory = defaultdict(list) # 奖励轨迹
        self.reward_item_names = [self.__class__.__name__] # 列表包含奖励函数的名称
        
    def reset(self, task, env):
        """
        这个方法用于在每个 episode 重置后执行奖励函数特定的重置操作。子类可以根据需要覆盖此方法来执行自定义的重置操作。
        如果 is_potential 为 True,则清除 pre_rewards 字典，并计算每个代理的初始奖励。
        """
        if self.is_potential:
            self.pre_rewards.clear()
            for agent_id in env.agents.keys():
                self.pre_rewards[agent_id] = self.get_reward(task, env, agent_id)
        self.reward_trajectory.clear()
        
    @abstractmethod
    def get_reward(self, task, env, agent_id):
        """
        这个方法是一个抽象方法，必须在子类中实现。它计算在当前时间步的奖励，并返回一个浮点数表示奖励的值。
        """
        raise NotImplementedError
    
    def _process(self, new_reward, agent_id, render_items=()):
        """
        这个方法用于处理奖励和内部变量。它接受一个新奖励值 new_reward,代理的 ID agent_id,以及可选的渲染项 render_items。
        它对奖励进行缩放，并在必要时计算潜在奖励。然后，将奖励和渲染项添加到奖励轨迹中，并返回奖励值。"""
        reward = new_reward * self.reward_scale
        if self.is_potential:
            reward, self.pre_rewards[agent_id] = reward - self.pre_rewards[agent_id], reward
        self.reward_trajectory[agent_id].append([reward, *render_items])
        return reward
    
    def get_reward_trajctory(self):
        """
        获取当前 episode 中的所有奖励历史记录。它返回一个字典，其中键是奖励函数的名称，值是奖励轨迹的 NumPy 数组
        """
        return dict(zip(self.reward_item_names, np.array(self.reward_trajectory.values()).transpose(2, 0, 1)))
            

        